{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "b31c0a84",
            "metadata": {},
            "source": [
                "## Rerun imports and initialization"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "1076c3a0",
            "metadata": {},
            "outputs": [],
            "source": [
                "from __future__ import annotations\n",
                "\n",
                "import math\n",
                "from collections import namedtuple\n",
                "from math import cos, sin\n",
                "\n",
                "import numpy as np\n",
                "import rerun as rr  # pip install rerun-sdk\n",
                "import rerun.blueprint as rrb\n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "bf894a1f",
            "metadata": {},
            "source": [
                "## Helper to create the colored cube\n",
                "\n",
                "This code exists in the `rerun.utilities` package, but is included here for context."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "f709925e",
            "metadata": {
                "jupyter": {
                    "source_hidden": true
                }
            },
            "outputs": [],
            "source": [
                "ColorGrid = namedtuple(\"ColorGrid\", [\"positions\", \"colors\"])\n",
                "\n",
                "\n",
                "def build_color_grid(x_count=10, y_count=10, z_count=10, twist=0):\n",
                "    \"\"\"\n",
                "    Create a cube of points with colors.\n",
                "\n",
                "    The total point cloud will have x_count * y_count * z_count points.\n",
                "\n",
                "    Parameters\n",
                "    ----------\n",
                "    x_count, y_count, z_count:\n",
                "        Number of points in each dimension.\n",
                "    twist:\n",
                "        Angle to twist from bottom to top of the cube\n",
                "\n",
                "    \"\"\"\n",
                "\n",
                "    grid = np.mgrid[\n",
                "        slice(-x_count, x_count, x_count * 1j),\n",
                "        slice(-y_count, y_count, y_count * 1j),\n",
                "        slice(-z_count, z_count, z_count * 1j),\n",
                "    ]\n",
                "\n",
                "    angle = np.linspace(-float(twist) / 2, float(twist) / 2, z_count)\n",
                "    for z in range(z_count):\n",
                "        xv, yv, zv = grid[:, :, :, z]\n",
                "        rot_xv = xv * cos(angle[z]) - yv * sin(angle[z])\n",
                "        rot_yv = xv * sin(angle[z]) + yv * cos(angle[z])\n",
                "        grid[:, :, :, z] = [rot_xv, rot_yv, zv]\n",
                "\n",
                "    positions = np.vstack([xyz.ravel() for xyz in grid])\n",
                "\n",
                "    colors = np.vstack(\n",
                "        [\n",
                "            xyz.ravel()\n",
                "            for xyz in np.mgrid[\n",
                "                slice(0, 255, x_count * 1j),\n",
                "                slice(0, 255, y_count * 1j),\n",
                "                slice(0, 255, z_count * 1j),\n",
                "            ]\n",
                "        ]\n",
                "    )\n",
                "\n",
                "    return ColorGrid(positions.T, colors.T.astype(np.uint8))\n"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "04c095ef",
            "metadata": {},
            "source": [
                "## Logging some data\n",
                "\n",
                "Now we can log some data and add it to the recording, and display it using `notebook_show`.\n",
                "\n",
                "Note that displaying a recording will consume the data, so it will not be available for use in a subsequent cell."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "92871ea1",
            "metadata": {},
            "outputs": [],
            "source": [
                "rr.init(\"rerun_example_cube\")\n",
                "\n",
                "STEPS = 100\n",
                "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
                "for t in range(STEPS):\n",
                "    rr.set_time_sequence(\"step\", t)\n",
                "    cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
                "    rr.log(\"cube\", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))\n",
                "\n",
                "rr.notebook_show()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "187430e1",
            "metadata": {},
            "source": [
                "## Logging live data\n",
                "\n",
                "Using `rr.notebook_show` like above buffers the data in the recording stream, but doesn't process it until the call to `rr.notebook_show`.\n",
                "\n",
                "However, `rr.notebook_show` can be called at any time during your cell's execution to immediately display the Rerun viewer. You can then incrementally stream to it. Here we add a sleep to simulate a cell that does a lot more processing. By calling `notebook_show` first, we can see the output of our code live while it's still running."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "246c8eb8",
            "metadata": {},
            "outputs": [],
            "source": [
                "from time import sleep\n",
                "\n",
                "rr.init(\"rerun_example_cube\")\n",
                "\n",
                "rr.notebook_show()\n",
                "\n",
                "STEPS = 100\n",
                "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
                "for t in range(STEPS):\n",
                "    sleep(0.05)\n",
                "    rr.set_time_sequence(\"step\", t)\n",
                "    cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
                "    rr.log(\"cube\", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "649b8e80-e69d-4b6b-be51-50cbf4c20495",
            "metadata": {},
            "source": [
                "## Incremental logging\n",
                "\n",
                "Note that until we either reset the recording stream (by calling `rr.init()`), or create a new output widget (by calling `rr.notebook_show()` The currently active stream in the kernel will continue to send events to the existing widget.\n",
                "\n",
                "The following will add a rotation to the above recording."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "7e4a9831-ac6c-4503-bcd7-b1212d0d542b",
            "metadata": {},
            "outputs": [],
            "source": [
                "for t in range(STEPS):\n",
                "    sleep(0.01)\n",
                "    rr.set_time_sequence(\"step\", t)\n",
                "    rr.log(\"cube\", rr.Transform3D(rotation=rr.RotationAxisAngle(axis=[1,0,0], degrees=t)))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "36f9f61b",
            "metadata": {},
            "source": [
                "## Starting a new recording\n",
                "\n",
                "You can always start another recording by calling `rr.init(...)` again to reset the global stream, or alternatively creating a separate recording stream using `rr.new_recording` (discussed more below)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "c4cc33fd",
            "metadata": {},
            "outputs": [],
            "source": [
                "rr.init(\"rerun_example_cube\")\n",
                "\n",
                "STEPS = 100\n",
                "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
                "for t in range(STEPS):\n",
                "    rr.set_time_sequence(\"step\", t)\n",
                "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
                "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
                "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
                "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))\n",
                "\n",
                "rr.notebook_show()\n"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "31d392a8",
            "metadata": {},
            "source": [
                "## Adjusting the view\n",
                "\n",
                "The  `notebook_show` method also lets you adjust properties such as width and height."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "1a1b0f66-4287-4705-8be5-ae837ffe3f90",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "rr.init(\"rerun_example_cube\")\n",
                "\n",
                "small_cube = build_color_grid(3, 3, 3, twist=0)\n",
                "rr.log(\"small_cube\", rr.Points3D(small_cube.positions, colors=small_cube.colors, radii=0.5))\n",
                "\n",
                "rr.notebook_show(width=400, height=400)\n"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "ff84c840",
            "metadata": {},
            "source": [
                "To update the default width and height, you can use the `rerun.notebook.set_default_size` function.\n",
                "\n",
                "Note that you do not need to initialize a recording to use this function."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "81157021",
            "metadata": {},
            "outputs": [],
            "source": [
                "from rerun.notebook import set_default_size\n",
                "\n",
                "set_default_size(width=400, height=400)\n",
                "\n",
                "rr.init(\"rerun_example_cube\")\n",
                "\n",
                "small_cube = build_color_grid(3, 3, 3, twist=0)\n",
                "rr.log(\"small_cube\", rr.Points3D(small_cube.positions, colors=small_cube.colors, radii=0.5))\n",
                "\n",
                "rr.notebook_show()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "755957e0",
            "metadata": {},
            "outputs": [],
            "source": [
                "set_default_size(width=640, height=480)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a9812634-067f-4e07-95fb-cb9a506c42d3",
            "metadata": {},
            "source": [
                "## Using blueprints\n",
                "\n",
                "Rerun blueprints can be used with `rr.notebook_show()`\n",
                "\n",
                "For example, we can split the two grids into their own respective views."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "eb8f7701",
            "metadata": {},
            "outputs": [],
            "source": [
                "rr.init(\"rerun_example_cube\")\n",
                "\n",
                "blueprint = rrb.Blueprint(\n",
                "    rrb.Horizontal(\n",
                "        rrb.Spatial3DView(name=\"Horizontal grid\", origin=\"h_grid\"),\n",
                "        rrb.Spatial3DView(name=\"Vertical grid\", origin=\"v_grid\"),\n",
                "        column_shares=[2,1]),\n",
                "    collapse_panels=True\n",
                ")\n",
                "\n",
                "rr.notebook_show(blueprint=blueprint)\n",
                "\n",
                "STEPS = 100\n",
                "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
                "for t in range(STEPS):\n",
                "    rr.set_time_sequence(\"step\", t)\n",
                "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
                "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
                "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
                "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "a63b30d7",
            "metadata": {},
            "source": [
                "## Extra convenience\n",
                "\n",
                "Rerun blueprints types also implement `_ipython_display_()` directly, so if a blueprint is the last element in your cell the right thing will happen.\n",
                "\n",
                "Note that this mechanism only works when you are using the global recording stream."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "6d1bf2ab",
            "metadata": {},
            "outputs": [],
            "source": [
                "rr.init(\"rerun_example_cube\")\n",
                "\n",
                "STEPS = 100\n",
                "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
                "for t in range(STEPS):\n",
                "    rr.set_time_sequence(\"step\", t)\n",
                "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
                "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
                "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
                "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))\n",
                "\n",
                "rrb.Spatial3DView(name=\"Horizontal grid\", origin=\"h_grid\")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "ef9087de-d090-4c90-ab3e-18c20c92bff6",
            "metadata": {},
            "source": [
                "## Working with non-global streams\n",
                "\n",
                "Sometimes it can be more explicit to work with specific (non-global recording) streams via the `new_recording` method.\n",
                "\n",
                "In this case, remember to call `notebook_show` directly on the recording stream. As noted above, there is no way to use a bare Blueprint object in conjunction with a non-global recording."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "ed86cc19-45cf-4c21-9a94-c4ce2ade7f65",
            "metadata": {},
            "outputs": [],
            "source": [
                "rec = rr.new_recording(\"rerun_example_cube_flat\")\n",
                "\n",
                "bp = rrb.Blueprint(collapse_panels=True)\n",
                "\n",
                "rec.notebook_show(blueprint=bp)\n",
                "\n",
                "flat_grid = build_color_grid(20, 20, 1, twist=0)\n",
                "rec.log(\"flat_grid\", rr.Points3D(flat_grid.positions, colors=flat_grid.colors, radii=0.5))"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.11.9"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}
